import torch
import numpy as np
import matplotlib.pyplot as plt


# 1. 收集数据
x = torch.rand([500])
y = 3 * x + 0.8

w = torch.rand(1, requires_grad=True)
b = torch.rand(1, requires_grad=True)

for i in range(5000):
    # 2. 开始训练
    y_pre = w * x + b

    # 3. 梯度置零
    for j in [w, b]:
        if j.grad:
            j.grad.data.zero_()

    # 4. 计算损失
    loss = (y_pre - y).pow(2).mean()

    # 5. 反向传播
    loss.backward()

    # 6. 更新参数
    w.data -= 0.01 * w.grad.data
    b.data -= 0.01 * b.grad.data

    if i % 500 == 0:
        print(w, b.data, loss.detach())

# 6. 可视化
y_predict = w * x + b
plt.style.use("seaborn-darkgrid")
plt.rc("font", size=20)
plt.rc("figure", figsize=(20, 8), dpi=100)
plt.scatter(x.numpy(), y.data.numpy(), c='r')
plt.plot(x.data.numpy(), y_predict.data.numpy())
plt.show()
print(y_predict.detach().numpy())


